package org.zlui65.milan.processor

import com.google.auto.service.AutoService
import com.squareup.kotlinpoet.FileSpec
import com.squareup.kotlinpoet.FunSpec
import com.squareup.kotlinpoet.TypeSpec
import java.io.File
import javax.annotation.processing.AbstractProcessor
import javax.annotation.processing.Processor
import javax.annotation.processing.ProcessingEnvironment
import javax.annotation.processing.RoundEnvironment
import javax.annotation.processing.SupportedAnnotationTypes
import javax.annotation.processing.SupportedSourceVersion
import javax.lang.model.SourceVersion
import javax.lang.model.element.Element
import javax.lang.model.element.ElementKind
import javax.lang.model.element.TypeElement
import javax.tools.Diagnostic

/**
 * Annotation processor for generating style extension functions.
 * This processor handles @Styleable, @StyleableChild, and @Style annotations.
 */
@AutoService(Processor::class)
@SupportedAnnotationTypes(
    "org.zlui65.milan.annotations.Styleable",
    "org.zlui65.milan.annotations.StyleableChild",
    "org.zlui65.milan.annotations.Style"
)
@SupportedSourceVersion(SourceVersion.RELEASE_17)
class StyleProcessor : AbstractProcessor() {

    override fun process(annotations: Set<TypeElement>, roundEnv: RoundEnvironment): Boolean {
        if (annotations.isEmpty()) return false

        // Process @Styleable classes
        val styleableElements = roundEnv.getElementsAnnotatedWith(
            try {
                Class.forName("org.zlui65.milan.annotations.Styleable")
            } catch (e: ClassNotFoundException) {
                processingEnv.messager.printMessage(
                    Diagnostic.Kind.WARNING,
                    "Could not find Styleable annotation class"
                )
                return false
            }
        )
        
        for (element in styleableElements) {
            if (element.kind != ElementKind.CLASS) {
                processingEnv.messager.printMessage(
                    Diagnostic.Kind.ERROR,
                    "@Styleable can only be applied to classes",
                    element
                )
                continue
            }

            processStyleableClass(element as TypeElement)
        }

        // Process @Style fields
        val styleElements = roundEnv.getElementsAnnotatedWith(
            try {
                Class.forName("org.zlui65.milan.annotations.Style")
            } catch (e: ClassNotFoundException) {
                processingEnv.messager.printMessage(
                    Diagnostic.Kind.WARNING,
                    "Could not find Style annotation class"
                )
                return false
            }
        )
        
        for (element in styleElements) {
            if (element.kind != ElementKind.FIELD) {
                processingEnv.messager.printMessage(
                    Diagnostic.Kind.ERROR,
                    "@Style can only be applied to fields",
                    element
                )
                continue
            }

            processStyleField(element)
        }

        return true
    }

    private fun processStyleableClass(element: TypeElement) {
        val packageName = processingEnv.elementUtils.getPackageOf(element).toString()
        val className = element.simpleName.toString()
        
        // Get the styleable name from the annotation
        val styleableName = try {
            val annotation = element.getAnnotation(Class.forName("org.zlui65.milan.annotations.Styleable"))
            val method = annotation.javaClass.getMethod("value")
            method.invoke(annotation) as String
        } catch (e: Exception) {
            processingEnv.messager.printMessage(
                Diagnostic.Kind.ERROR,
                "Failed to get styleable name from annotation: ${e.message}",
                element
            )
            return
        }

        // Generate extension functions for this styleable class
        generateStyleableExtensions(packageName, className, styleableName, element)
    }

    private fun processStyleField(element: Element) {
        val enclosingElement = element.enclosingElement as TypeElement
        val packageName = processingEnv.elementUtils.getPackageOf(enclosingElement).toString()
        val className = enclosingElement.simpleName.toString()
        val fieldName = element.simpleName.toString()

        // Generate extension functions for this style field
        generateStyleExtensions(packageName, className, fieldName, element)
    }

    private fun generateStyleableExtensions(
        packageName: String,
        className: String,
        styleableName: String,
        element: TypeElement
    ) {
        val fileName = "${className}StyleExtensions"
        
        // Create a new file using KotlinPoet
        val fileBuilder = FileSpec.builder(packageName, fileName)
            .addComment("Generated code for $className style extensions")
            .addImport("org.zlui65.milan", "Style", "StyleBuilder", "style")
        
        // Find all @StyleableChild fields
        val styleableChildFields = element.enclosedElements.filter { 
            it.kind == ElementKind.FIELD && 
            it.getAnnotationsByType(Class.forName("org.zlui65.milan.annotations.StyleableChild")).isNotEmpty()
        }
        
        for (field in styleableChildFields) {
            val fieldName = field.simpleName.toString()
            
            // Get the attribute ID from the annotation
            val attrId = try {
                val annotation = field.getAnnotation(Class.forName("org.zlui65.milan.annotations.StyleableChild"))
                val method = annotation.javaClass.getMethod("value")
                method.invoke(annotation) as Int
            } catch (e: Exception) {
                processingEnv.messager.printMessage(
                    Diagnostic.Kind.ERROR,
                    "Failed to get attribute ID from annotation: ${e.message}",
                    field
                )
                continue
            }
            
            // Generate extension function for setting this child's style
            val styleResFun = FunSpec.builder("${fieldName}Style")
                .addParameter("styleRes", Int::class)
                .addStatement("setAttribute($attrId, styleRes)")
                .receiver(Class.forName("org.zlui65.milan.StyleBuilder"))
                .build()
            
            // Generate extension function for building this child's style
            val styleBuilderFun = FunSpec.builder("${fieldName}Style")
                .addParameter("builder", Class.forName("kotlin.Function1"))
                .addStatement("val styleBuilder = %T()", Class.forName("org.zlui65.milan.StyleBuilder"))
                .addStatement("styleBuilder.builder()")
                .addStatement("setAttribute($attrId, styleBuilder.build())")
                .receiver(Class.forName("org.zlui65.milan.StyleBuilder"))
                .build()
            
            fileBuilder.addFunction(styleResFun)
            fileBuilder.addFunction(styleBuilderFun)
        }
        
        // Write the generated file
        try {
            val kaptKotlinGeneratedDir = processingEnv.options["kapt.kotlin.generated"]
                ?: return
            
            val file = File(kaptKotlinGeneratedDir)
            fileBuilder.build().writeTo(file)
        } catch (e: Exception) {
            processingEnv.messager.printMessage(
                Diagnostic.Kind.ERROR,
                "Failed to write generated file: ${e.message}"
            )
        }
    }

    private fun generateStyleExtensions(
        packageName: String,
        className: String,
        fieldName: String,
        element: Element
    ) {
        val fileName = "${className}StyleExtensions"
        
        // Create a new file using KotlinPoet
        val fileBuilder = FileSpec.builder(packageName, fileName)
            .addComment("Generated code for $className style extensions")
            .addImport("org.zlui65.milan", "Style", "StyleBuilder", "style")
        
        // Generate extension function for adding this style
        val addStyleFun = FunSpec.builder("add${fieldName.capitalize()}")
            .addStatement("add($className.$fieldName)")
            .receiver(Class.forName("org.zlui65.milan.StyleBuilder"))
            .build()
        
        fileBuilder.addFunction(addStyleFun)
        
        // Write the generated file
        try {
            val kaptKotlinGeneratedDir = processingEnv.options["kapt.kotlin.generated"]
                ?: return
            
            val file = File(kaptKotlinGeneratedDir)
            fileBuilder.build().writeTo(file)
        } catch (e: Exception) {
            processingEnv.messager.printMessage(
                Diagnostic.Kind.ERROR,
                "Failed to write generated file: ${e.message}"
            )
        }
    }

    private fun String.capitalize(): String {
        return if (this.isEmpty()) this
        else this[0].uppercaseChar() + this.substring(1)
    }
}
